package com.xw.llm;

import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.WebFluxSseClientTransport;
import io.modelcontextprotocol.spec.McpSchema;
import org.springframework.web.reactive.function.client.WebClient;

import java.time.Duration;
import java.util.List;
import java.util.Map;

public class McpClientTest {

    /**
     * https://github.com/modelcontextprotocol/java-sdk/blob/0.10.x/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java
     * https://mcpcn.com/docs/sdk/java/mcp-client/
     *
     * @param args
     */
    public static void main(String[] args) {
        WebClient.Builder webClientBuilder = null;
        WebFluxSseClientTransport transport = null;
        McpSchema.CallToolRequest callToolRequest = null;

        int type = 2;

        if (type == 1) {
            webClientBuilder = WebClient.builder().baseUrl("https://mcp.api-inference.modelscope.net");
            transport = WebFluxSseClientTransport.builder(webClientBuilder).sseEndpoint("/183f13cb68d54b/sse").build();
            callToolRequest = new McpSchema.CallToolRequest("bing_search", Map.of("query", "今天上海天气"));
        } else if (type == 2) {
            webClientBuilder = WebClient.builder().baseUrl("http://localhost:8080");
            transport = WebFluxSseClientTransport.builder(webClientBuilder).sseEndpoint("/sse").build();
            callToolRequest = new McpSchema.CallToolRequest("getOrderDetail", Map.of("query", "今天上海天气"));
        }


        McpSyncClient client = McpClient.sync(transport)
                .requestTimeout(Duration.ofSeconds(100))
                .loggingConsumer(notification -> {
                    System.out.println("Received log message: " + notification.data());
                })
                .capabilities(McpSchema.ClientCapabilities.builder()
                        .roots(true)      // Enable roots capability
                        .sampling()       // Enable sampling capability
                        .build())
                .build();

        McpSchema.InitializeResult initialize = client.initialize();

        Object ping = client.ping();

        McpSchema.ListToolsResult listToolsResult = client.listTools();

        McpSchema.CallToolResult callToolResult = client.callTool(callToolRequest);

        List<McpSchema.Content> content = callToolResult.content();
        System.out.println();
    }
}
